"""
完全重置系统数据
清空数据库所有表 + 删除所有上传文件和向量数据
⚠️ 警告：此操作不可逆！
"""
import sys
import os
import shutil
from pathlib import Path

# 添加项目根目录到 Python 路径
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from sqlalchemy import text, create_engine
from app.core.config import settings
from loguru import logger

logger.add("logs/reset_data.log", rotation="1 day", level="INFO")


def confirm_reset():
    """确认是否真的要重置"""
    print("=" * 80)
    print("⚠️  警告：即将清空所有数据！")
    print("=" * 80)
    print()
    print("此操作将：")
    print("  1. 清空数据库所有表（保留表结构）")
    print("  2. 删除所有上传的文件")
    print("  3. 删除所有向量数据")
    print()
    print("⚠️  此操作不可逆！")
    print()
    
    response = input("确认要继续吗？输入 'YES' 继续，其他任何输入取消: ")
    return response == "YES"


def clear_database():
    """清空数据库所有表"""
    logger.info("开始清空数据库...")
    print("\n[1/3] 清空数据库...")
    
    try:
        engine = create_engine(settings.DATABASE_URL)
        
        with engine.connect() as conn:
            # 获取所有表
            tables_query = text("""
                SELECT tablename 
                FROM pg_tables 
                WHERE schemaname = 'public'
            """)
            result = conn.execute(tables_query)
            tables = [row[0] for row in result]
            
            print(f"  找到 {len(tables)} 个表")
            
            # 清空每个表
            for table in tables:
                try:
                    # 使用 TRUNCATE CASCADE 清空表及其关联数据
                    truncate_query = text(f"TRUNCATE TABLE {table} RESTART IDENTITY CASCADE")
                    conn.execute(truncate_query)
                    print(f"  ✓ 清空表: {table}")
                    logger.info(f"清空表: {table}")
                except Exception as e:
                    print(f"  ✗ 清空表 {table} 失败: {e}")
                    logger.error(f"清空表 {table} 失败: {e}")
            
            conn.commit()
            
        print("  ✅ 数据库清空完成")
        logger.info("数据库清空完成")
        return True
        
    except Exception as e:
        print(f"  ❌ 数据库清空失败: {e}")
        logger.error(f"数据库清空失败: {e}")
        return False


def clear_uploads():
    """删除所有上传的文件"""
    logger.info("开始删除上传文件...")
    print("\n[2/3] 删除上传文件...")
    
    try:
        upload_dir = Path(settings.UPLOAD_DIR)
        
        if upload_dir.exists():
            # 统计文件数量
            file_count = sum(1 for _ in upload_dir.rglob('*') if _.is_file())
            print(f"  找到 {file_count} 个文件")
            
            # 删除整个目录
            shutil.rmtree(upload_dir)
            print(f"  ✓ 删除目录: {upload_dir}")
            
            # 重新创建空目录
            upload_dir.mkdir(parents=True, exist_ok=True)
            print(f"  ✓ 重新创建空目录")
            
            print("  ✅ 上传文件清理完成")
            logger.info(f"删除了 {file_count} 个上传文件")
        else:
            print("  ℹ️  上传目录不存在，跳过")
            logger.info("上传目录不存在")
        
        return True
        
    except Exception as e:
        print(f"  ❌ 删除上传文件失败: {e}")
        logger.error(f"删除上传文件失败: {e}")
        return False


def clear_vector_stores():
    """删除所有向量数据"""
    logger.info("开始删除向量数据...")
    print("\n[3/3] 删除向量数据...")
    
    try:
        vector_dir = Path(settings.VECTOR_STORE_DIR)
        
        if vector_dir.exists():
            # 统计向量库数量
            vector_count = sum(1 for _ in vector_dir.iterdir() if _.is_dir())
            print(f"  找到 {vector_count} 个向量库")
            
            # 删除整个目录
            shutil.rmtree(vector_dir)
            print(f"  ✓ 删除目录: {vector_dir}")
            
            # 重新创建空目录
            vector_dir.mkdir(parents=True, exist_ok=True)
            print(f"  ✓ 重新创建空目录")
            
            print("  ✅ 向量数据清理完成")
            logger.info(f"删除了 {vector_count} 个向量库")
        else:
            print("  ℹ️  向量目录不存在，跳过")
            logger.info("向量目录不存在")
        
        return True
        
    except Exception as e:
        print(f"  ❌ 删除向量数据失败: {e}")
        logger.error(f"删除向量数据失败: {e}")
        return False


def recreate_admin():
    """重新创建管理员账户"""
    logger.info("开始创建管理员账户...")
    print("\n[4/4] 创建管理员账户...")
    
    try:
        import uuid
        from sqlalchemy.orm import Session
        from app.core.security import get_password_hash
        from app.models.database import engine
        
        # 先导入所有模型以避免关系映射错误
        from app.models.user import User, UserRole
        from app.models.document import Document
        from app.models.document_share import DocumentShare
        from app.models.category import Category
        from app.models.knowledge_base import KnowledgeBase
        from app.models.chat_history import ChatSession, ChatMessage
        
        with Session(engine) as session:
            admin_user = User(
                id=str(uuid.uuid4()),
                username="admin",
                email="admin@shihe.com",
                hashed_password=get_password_hash("admin123"),
                full_name="系统管理员",
                role=UserRole.ADMIN,
                is_active=True
            )
            session.add(admin_user)
            session.commit()
            
            print("  ✅ 管理员账户创建成功")
            print("     用户名: admin")
            print("     密码: admin123")
            logger.info("管理员账户创建成功")
        
        return True
        
    except Exception as e:
        print(f"  ❌ 创建管理员账户失败: {e}")
        logger.error(f"创建管理员账户失败: {e}")
        return False


def main():
    """主函数"""
    print("\n")
    print("=" * 80)
    print("  思和智库 - 数据重置工具")
    print("=" * 80)
    print()
    
    # 确认操作
    if not confirm_reset():
        print("\n❌ 操作已取消")
        logger.info("用户取消了重置操作")
        return
    
    print("\n开始重置...")
    logger.info("=" * 60)
    logger.info("开始重置系统数据")
    logger.info("=" * 60)
    
    # 执行清理
    success = True
    
    # 1. 清空数据库
    if not clear_database():
        success = False
    
    # 2. 删除上传文件
    if not clear_uploads():
        success = False
    
    # 3. 删除向量数据
    if not clear_vector_stores():
        success = False
    
    # 4. 重新创建管理员
    if success:
        recreate_admin()
    
    # 总结
    print("\n" + "=" * 80)
    if success:
        print("✅ 数据重置完成！")
        print()
        print("系统已恢复到初始状态：")
        print("  • 数据库：所有表已清空")
        print("  • 文件：所有上传文件已删除")
        print("  • 向量：所有向量数据已删除")
        print("  • 用户：已创建默认管理员账户")
        print()
        print("默认管理员账户：")
        print("  用户名: admin")
        print("  密码: admin123")
        print()
        print("⚠️  请重启后端服务使更改生效")
        logger.info("数据重置成功完成")
    else:
        print("❌ 数据重置过程中出现错误")
        print("请查看日志文件: logs/reset_data.log")
        logger.error("数据重置过程中出现错误")
    print("=" * 80)
    print()


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n\n❌ 操作被用户中断")
        logger.info("操作被用户中断")
    except Exception as e:
        print(f"\n\n❌ 发生未预期的错误: {e}")
        logger.error(f"发生未预期的错误: {e}")
        import traceback
        logger.error(traceback.format_exc())

